package com.long1008.common;

import org.dmg.pmml.FieldName;
import org.dmg.pmml.PMML;
import org.jpmml.evaluator.Evaluator;
import org.jpmml.evaluator.InputField;
import org.jpmml.evaluator.ModelEvaluatorFactory;
import org.jpmml.model.PMMLUtil;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;
import org.xml.sax.SAXException;

import javax.xml.bind.JAXBException;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Service(value = "dTreeUtils")
public class DTreeUtils {
    private Evaluator modelEvaluator;

    /**
     * 加载模型
     * @param pmmlFileName
     * @return
     */
    public Evaluator LoadPmml(String pmmlFileName){
        PMML pmml= new PMML();
        try {
            if(pmmlFileName != null){
                InputStream is = new FileInputStream(pmmlFileName);
                pmml = PMMLUtil.unmarshal(is);
                try {
                    is.close();
                }catch (IOException e){
                    System.out.println("InputStream close error!");
                }
            }
            ModelEvaluatorFactory modelEvaluatorFactory = ModelEvaluatorFactory.newInstance();
            this.modelEvaluator = (Evaluator) modelEvaluatorFactory.newModelEvaluator(pmml);

            modelEvaluator.verify();
            System.out.println("加载模型成功！");



        }catch (SAXException e){
            e.printStackTrace();
        }catch (JAXBException e){
            e.printStackTrace();
        }catch (FileNotFoundException e){
            e.printStackTrace();
        }

        return modelEvaluator;
    }

    /**
     * 获取模型所有特征名
     * @param modelEvaluator
     * @return
     */
    public List<String> getFeatureName(Evaluator modelEvaluator){
        this.modelEvaluator = modelEvaluator;

        List<String> featureNames = new ArrayList<>();
        List<InputField> inputFields = modelEvaluator.getInputFields();

        for (InputField inputField:inputFields){
            featureNames.add(inputField.getName().toString());
        }

        System.out.println(featureNames);
        return featureNames;
    }

    /**
     * 获取模型目标字段名
     * @param modelEvaluator
     * @return
     */
    public FieldName getTargetName(Evaluator modelEvaluator){
        this.modelEvaluator = modelEvaluator;
        return modelEvaluator.getTargetFields().get(0).getName();
    }

    /**
     * 预测模型获取目标值
     * @param modelEvaluator
     * @param arguments
     * @return
     */
    public Object predict(Evaluator modelEvaluator, Map<FieldName, Object> arguments){
        this.modelEvaluator = modelEvaluator;
        Map<FieldName, ?> result = modelEvaluator.evaluate(arguments);

        FieldName TargetName = getTargetName(modelEvaluator);

        Object targetValue = result.get(TargetName);
        return targetValue;
    }

}
